from ui import ImageProcessorUI
from .base import ImageProcessor
from PIL import Image
import numpy as np
import cv2

class CellCountProcessor(ImageProcessor):
    def process(self, image: Image.Image, **kwargs) -> Image.Image:
        """
        对图像中的细胞进行计数。
        参数:
        - image: 输入图像 (PIL Image)
        - min_cell_size: 细胞的最小尺寸 (默认为100像素)
        返回:
        - Image.Image: 包含计数结果的处理后图像
        """
        min_cell_size = kwargs.get("min_cell_size", 100)
        # 转换为灰度图
        img_array = np.array(image.convert('L'))
        # 自适应直方图均衡化（CLAHE）以增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        img_clahe = clahe.apply(img_array)
        # 高斯模糊减少噪声
        img_blur = cv2.GaussianBlur(img_clahe, (5, 5), 0)
        # 使用自适应阈值分割
        thresh = cv2.adaptiveThreshold(
            img_blur,
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY_INV,  # 反色二值化，细胞为白色
            blockSize=15,
            C=3
        )
        # 形态学操作：开运算去除小噪点，闭运算连接断裂区域
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
        # 连通域分析
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(thresh, connectivity=8)
        # 创建彩色图像用于绘制结果
        result_image = cv2.cvtColor(img_array, cv2.COLOR_GRAY2BGR)
        cell_count = 0
        for i in range(1, num_labels):  # 跳过背景
            area = stats[i, cv2.CC_STAT_AREA]
            if area >= min_cell_size:
                x, y, w, h = stats[i, cv2.CC_STAT_LEFT], stats[i, cv2.CC_STAT_TOP], \
                    stats[i, cv2.CC_STAT_WIDTH], stats[i, cv2.CC_STAT_HEIGHT]
                # 计算圆形度，排除长条形伪影
                aspect_ratio = float(w) / h if h > 0 else 0
                if 0.5 < aspect_ratio < 2.0:  # 假设细胞近似圆形
                    cv2.rectangle(result_image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 0), 2)
                    cell_count += 1
        # 显示计数结果
        cv2.putText(
            result_image,
            f'Cells: {cell_count}',
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (255, 0, 0),
            2,
            cv2.LINE_AA
        )
        return Image.fromarray(result_image)

    def get_ui_parameters(self, ui: "ImageProcessorUI") -> dict:
        """
        从 UI 获取参数
        """
        try:
            return {
                "min_cell_size": 100
            }
        except Exception as e:
            ui.show_error(f"获取参数失败: {e}")
            return {}